# -*- coding: utf-8 -*-

"""
Created on Fri Jun  8 18:43:24 2018

@author: Administrator
"""

import os
import matplotlib as mpl
#这一行代码是指保存图片，不在console里面显示。
mpl.use('Agg') 
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
from  scipy.stats import pearsonr 
#%%
        
#%%
class Basics(object):
    underline = '_'
    linkline = '-'
    style = 'ticks'
    context = 'paper'
    palette = 'PRGn'
    svg = 'svg'
    dot_svg = '.svg'
#%%

class Boxplot(Basics):
    #作用原理：
    #参数说明：
    #参数说明：
    #返回值说明：
    #举例：
    #调用：
    #被调用：
    #bug：
    def __init__(self,coeff_pvalue_df,dvar,idvar_list,output_filepath):
        print('in the Boxplot')
        #作图
        sns.set(style=self.style)
        sns.set_context(self.context)
        boxplot = sns.boxplot(y=dvar,x=idvar_list,data=coeff_pvalue_df,palette=self.palette)
        sns.despine(offset=10,trim=True)
        boxplot.get_figure().savefig(output_filepath)
        boxplot.clear() #this is import when plot 2 graph at the same time.


#%%
class Barplot(Basics):

#    #作用原理：
#    #参数说明：
#    #参数说明：
#    #返回值说明：
#    #举例：
#    #调用：
#    #被调用：
#    #bug：
    def __init__(self,coeff_pvalue_df,column_name_list,x,y,output_dirpath,output_filename):
        print('in the Barplot')
        boxplot_file_str = os.sep + r'across_group_barplot_'
        #作图
        sns.set(style=self.style)
        sns.set_context(self.context)
        barplot = sns.barplot(x=column_name_list[x],y=column_name_list[y],data=coeff_pvalue_df,palette=self.palette)
        sns.despine(offset=10,trim=True)
        barplot.get_figure().savefig(output_dirpath+boxplot_file_str+output_filename)
        barplot.clear() #this is import when plot 2 graph at the same time.
               
#%%
        
class Juxtaposed_barplot(Basics):

#    #作用原理：
#    #参数说明：list0,list1是需要做柱状图的两列数，list0应为大的那一列。
#    #参数说明：label_list,分别对应两列数的名称。
#    #返回值说明：
#    #举例：
#    #调用：
#    #被调用： 
#    #bug：
    barwidth = 0.6
    color_code0 = '#1478B4'  #深蓝'#002060'
    color_code1 = '#FA7F14'  #浅蓝'#5B9BD5'
    subj = 'subjectID'
    step = 5
    def __init__(self,list0,list1,y_label,legend_list,output_filepath):
        print('in the Barplot')
        print('list0:',list0)
        print('list1:',list1)
        numb0,numb1 = len(list0),len(list1)
        if not numb0 == numb1:
          #column_name_list对应的df中列的元素数目必须相等，否则报错
          raise Exception('df[column_name_list[0]] and df[column_name_list[1]] must have same number of elements' + str(numb0) + '  ' + str(numb1))
        
        x_coordinate0 = np.arange(numb1) * 3 #确定bar在x上的坐标
        
        medi = np.arange(1,numb1,step=5)
        tick_label_list = [str(i) if i in medi else '' for i in np.arange(numb1) + 1]

        #作图
        plt.style.use('ggplot') #使用ggplot的特点作图
        fig = plt.figure(figsize=(6,3))
        ax = fig.add_subplot(111)
        ax.bar(x_coordinate0, list0, width=self.barwidth, label=legend_list[0],color=self.color_code0)
        ax.bar(x_coordinate0+self.barwidth, list1, width=self.barwidth, label=legend_list[1],\
               color=self.color_code1)
        ax.set_xticks(x_coordinate0 + self.barwidth/2)
        ax.set_xticklabels(labels=tick_label_list)
        plt.xlabel(self.subj)
        plt.ylabel(y_label)
        ax.tick_params(length=0)
        plt.legend()
        plt.tight_layout(pad=0.5)#调整图周围的空白
        fig.savefig(output_filepath,dpi=1000)


#%%
class Error_barplot(Basics,Axe_bar,Special_division):
  one_numb = 1  #用于array.reshape()
  zero_numb = 0 #用于从ax_arrayp[][]取出一个ax
  x_label_str = ''
  y_label_str = 'Correlation'
  bar_color_list = ['#539caf']
  error_color_list = ['#000000']
  def __init__(self,df,axe_column_name,bar_column_name,interest_column_name,figsize_tuple,column_numb,barwidth,spacing):

      #下面开始构造figure,使用ggplot的特点作图
      plt.style.use('ggplot')
      self.fig = plt.figure(figsize=figsize_tuple)
      axe_list = df[axe_column_name].unique()
      #通过axe_list的个数，来得出Figure有几个axe
      #如果len(axe_list)可以被self.column_numb整除，row_numb等于取整的商，如果不能，则等于取整的商+1
      row_numb = self.special_division(len(axe_list),column_numb)
      #sharex=True, sharey=True还暂不确定是什么意思,之所以加reshape(),是为了将二维array拉成一维的，
      #便于for循环。
      ax_array = self.fig.subplots(row_numb,column_numb, sharex=True, sharey=True).reshape(row_numb*column_numb,self.one_numb)
      
      bar_list = df[bar_column_name].unique()
      self.mean_std_df = pd.DataFrame(index=bar_list,columns=axe_list)     

      #下面每一个axe_item占一个axe，每一个bar_item占某一个axe内的一个bar

      i = 0
      for axe_item in axe_list:
        #下面将df属于axe_item的那一部分取出
        axe_item_df = df[df[axe_column_name]==axe_item]
        x_data = np.arange(len(bar_list)) * spacing
        y_data = []  #需要在循环中给予赋值
        error_data = []   #需要在循环中给予赋值
        x_label = self.x_label_str
        y_label = self.y_label_str
        title = axe_item
        #下面开始按照bar_item循环，以收集y_data和error_data
        j = 0
        for bar_item in bar_list:
          #agent是一个中间变量
          agent = axe_item_df[axe_item_df[bar_column_name]==bar_item][interest_column_name]
          mean = '%.2f'%agent.mean()
          std = '%.2f'%agent.std()
          y_data.append(float(mean))
          error_data.append(float(std))
          
          self.mean_std_df[axe_item][bar_item] = mean + '±' + std
          j = j + 1

        self.axe(ax_array[i][self.zero_numb],x_data, y_data, error_data, x_label, y_label, title,\
                 self.bar_color_list[0],barwidth,self.error_color_list[0],bar_list)
        i = i + 1

#%%

class Error_barplot_across_region(Basics,Axe_bar):
  zero_numb = 0 #用于从ax_arrayp[][]取出一个ax
  x_label_str = ''
  y_label_str = 'Correlation'
  title_str = 'Across region correlation'
  bar_color_list = ['#539caf']
  error_color_list = ['#000000']
  def __init__(self,df,modality_rm_list,subj_column_name,modality_column_name,figsize_tuple,barwidth,spacing):

      #下面开始构造figure,使用ggplot的特点作图
      plt.style.use('ggplot')
      self.fig = plt.figure(figsize=figsize_tuple)
      #sharex=True, sharey=True还暂不确定是什么意思
      ax = self.fig.subplots(1,1,sharex=True, sharey=True)
      
      modality_list = df[modality_column_name].unique() 
      subj_list = df[subj_column_name].unique()
      x_data = np.arange(len(modality_list)) * spacing
      y_data = []  #需要在循环中给予赋值
      error_data = [] #需要在循环中给予赋值
      x_label = self.x_label_str
      y_label = self.y_label_str
      title = self.title_str
      coeff_df = pd.DataFrame(index=subj_list,columns=modality_list)
      for modality_item in modality_list:
        modality_item_df = df[df[modality_column_name]==modality_item]
        for subj_item in subj_list:
          subj_modality_item_df = modality_item_df[modality_item_df[subj_column_name]==subj_item]
          first_modality_rms = subj_modality_item_df[modality_rm_list[0]]
          second_modality_rms = subj_modality_item_df[modality_rm_list[1]]
          coeff_df[modality_item][subj_item] = pearsonr(first_modality_rms,second_modality_rms)
        x_data = '%.2f'%coeff_df[modality_item].mean()
        error_data = '%.2f'%coeff_df[modality_item].std()
      self.axe(ax,x_data, y_data, error_data, x_label, y_label, title,\
               self.bar_color_list[0],barwidth,self.error_color_list[0],modality_list)

#%%

class  Coeff_barplot_across_subject_regionwise(Basics,Axe_bar,Special_division):
  one_numb = 1  #用于array.reshape()
  zero_numb = 0 #用于从ax_arrayp[][]取出一个ax
  x_label_str = ''
  y_label_str = 'Correlation'
  bar_color_list = ['#539caf']
  error_color_list = ['#000000']
  def __init__(self,df,axe_column_name,bar_column_name,modality_rm_list,figsize_tuple,column_numb,barwidth,spacing):

      #下面开始构造figure,使用ggplot的特点作图
      plt.style.use('ggplot')
      self.fig = plt.figure(figsize=figsize_tuple)
      axe_list = df[axe_column_name].unique()
      #通过axe_list的个数，来得出Figure有几个axe
      #如果len(axe_list)可以被self.column_numb整除，row_numb等于取整的商，如果不能，则等于取整的商+1
      row_numb = self.special_division(len(axe_list),column_numb)
      #sharex=True, sharey=True还暂不确定是什么意思,之所以加reshape(),是为了将二维array拉成一维的，
      #便于for循环。
      ax_array = self.fig.subplots(row_numb,column_numb, sharex=True, sharey=True).reshape(row_numb*column_numb,self.one_numb)
      
      bar_list = df[bar_column_name].unique()
      self.coeff_across_subject_regionwise_df = pd.DataFrame(index=bar_list,columns=axe_list)

      #下面每一个axe_item占一个axe，每一个bar_item占某一个axe内的一个bar

      i = 0
      for axe_item in axe_list:
        #下面将df属于axe_item的那一部分取出
        axe_item_df = df[df[axe_column_name]==axe_item]
        x_data = np.arange(len(bar_list)) * spacing
        y_data = []  #需要在循环中给予赋值
        error_data = []   #需要在循环中给予赋值
        x_label = self.x_label_str
        y_label = self.y_label_str
        title = axe_item
        #下面开始按照bar_item循环，以收集y_data和error_data
        j = 0
        for bar_item in bar_list:
          #agent是一个中间变量
          modality1_rm_nparray = axe_item_df[axe_item_df[bar_column_name]==bar_item][modality_rm_list[0]]
          modality2_rm_nparray = axe_item_df[axe_item_df[bar_column_name]==bar_item][modality_rm_list[1]]
          #计算modality1_rm_nparray,modality2_rm_nparray这两个nparray的Pearson coefficient
          coeff = pearsonr(modality1_rm_nparray,modality2_rm_nparray)
          y_data.append(coeff)
          
          self.coeff_across_subject_regionwise_df[axe_item][bar_item] = coeff
          j = j + 1

        self.axe(ax_array[i][self.zero_numb],x_data, y_data, error_data, x_label, y_label, title,\
                 self.bar_color_list[0],barwidth,self.error_color_list[0],bar_list)
        i = i + 1

#%%